from ICF.FileLoader import FileLoader
from typing import Set, List, Tuple, Dict
from tqdm import tqdm
import pickle
import os


class Trainer:
    def __init__(self, hyper_params):
        self.file_loader = FileLoader(hyper_params)

        # Create Interaction Set
        print('ICF: Trainer: Loading Interaction List...')
        self.interaction_list: Dict[int, List[int]] = {}

        for user, history in self.file_loader.intereaction_list.items():
            self.interaction_list[user] = []
            for item in history:
                self.interaction_list[user].append(item.item_id)
        print('ICF: Trainer: Finished.')

        try:
            self.W = pickle.load(open('ICF/model_dat/W.pkl', 'rb'))
            print('ICF: Model Loaded.')
        except Exception as err:
            self.cal_similarity()

    def cal_similarity(self):
        # Co-rated items between users
        C = {}
        N = {}
        W = {}
        print('ICF: Calculating similarity between items...')
        for user, items in tqdm(self.interaction_list.items()):
            for i in items:
                if i not in N.keys():
                    N[i] = 0
                N[i] += 1
                if i not in C.keys():
                    C[i] = {}
                    W[i] = {}
                for j in items:
                    if j == i:
                        continue
                    if j not in C[i].keys():
                        C[i][j] = 0
                        W[i][j] = 0
                    C[i][j] += 1
        print('ICF: Co-rated items count finished')
        # Calculate similarity matrix
        for i, related_items in tqdm(C.items()):
            for j, cij in related_items.items():
                W[i][j] = cij/(N[i]*N[j])**0.5
        print('ICF: Similarity calculation finished')

        self.W = W
        os.makedirs('ICF/model_dat', exist_ok=True)
        pickle.dump(self.W, open('ICF/model_dat/W.pkl', 'wb'))


if __name__ == '__main__':
    hyper_params = {
        'dataset_path': '../datasets/lib.txt'
    }

    trainer = Trainer(hyper_params)
